package org.thanos.demo.limit;

import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ValueOperations;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.lang.reflect.Method;
import java.net.InetSocketAddress;
import java.nio.file.AccessDeniedException;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * 限流工具
 *
 * @author Administrator
 */
@Aspect
@Component
public class RateLimitAspect {

    private static final Logger LOGGER = LoggerFactory.getLogger(RateLimitAspect.class);

    private static final Integer RATE_LIMIT = 10;
    private static final Integer RATE_LIMIT_TIMEOUT = 30;

    @Autowired
    private StringRedisTemplate redisTemplate;

    @Pointcut("@annotation(org.thanos.demo.limit.RateLimit)")
    public void annotationPointCut() {
    }

    @Before("annotationPointCut()")
    public void doBefore(JoinPoint jp) throws IOException {
        Method sourceMethod = getSourceMethod(jp);
        if (Objects.nonNull(sourceMethod)) {
            RateLimit rateLimit = sourceMethod.getAnnotation(RateLimit.class);
            if (Objects.nonNull(rateLimit)) {
                int requestIndex = rateLimit.requestIndex();
                Object[] args = jp.getArgs();
                ServerHttpRequest request = (ServerHttpRequest) args[requestIndex];
                String clientIp = Objects.requireNonNull(request.getRemoteAddress()).getAddress().getHostAddress();
                ValueOperations<String, String> operations = redisTemplate.opsForValue();
                if (redisTemplate.hasKey(clientIp)) {
                    //这里是防止原子操作会把过期时间消除
                    if ("1".equals(operations.get(clientIp))) {
                        redisTemplate.expire(clientIp, RATE_LIMIT_TIMEOUT, TimeUnit.SECONDS);
                    }
                    String value = operations.get(clientIp);
                    if (!StringUtils.isEmpty(value) && Integer.parseInt(value) > RATE_LIMIT) {
                        throw new AccessDeniedException("操作过于频繁");
                    } else {
                        operations.increment(clientIp);
                    }
                } else {
                    operations.increment(clientIp);
                }
            }
        }

    }

    private Method getSourceMethod(JoinPoint jp) {
        Method proxyMethod = ((MethodSignature) jp.getSignature()).getMethod();
        try {
            return jp.getTarget().getClass().getMethod(proxyMethod.getName(), proxyMethod.getParameterTypes());
        } catch (NoSuchMethodException | SecurityException e) {
            LOGGER.error(e.getMessage(), e);
        }
        return null;
    }
}